#ifndef __WINDBGEXTIMPL_H__
#define __WINDBGEXTIMPL_H__

#include <string.h>
#include "../common/WindbgExt.h"

#ifndef IMAGE_FILE_MACHINE_ARMNT
#define IMAGE_FILE_MACHINE_ARMNT 0xff
#endif

#include <engextcpp.hpp>


//////////////////////////////////////////////////////////////////////////
// Version
#define WINDBGEXT_VERSION		0x00010000
#define WINDBGEXT_VERSION_STR	"1.0"


//////////////////////////////////////////////////////////////////////////
// CaptureOutput

class CaptureOutputAImpl: public CaptureOutputA
{
public:
	CaptureOutputAImpl(): m_ref(1)
	{
	}

	virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void **ppvObject)
	{
		return E_NOTIMPL;
	}

	virtual ULONG STDMETHODCALLTYPE AddRef(void)
	{
		return InterlockedIncrement(&m_ref);
	}

	virtual ULONG STDMETHODCALLTYPE Release(void)
	{
		ULONG ret = InterlockedDecrement(&m_ref);
		if (ret == 0)
			delete this;
		return ret;
	}

	virtual void STDMETHODCALLTYPE Start(void)
	{
		m_capture.Start();
	}

	virtual void STDMETHODCALLTYPE Stop(void)
	{
		m_capture.Stop();
	}

	virtual void STDMETHODCALLTYPE Delete(void)
	{
		m_capture.Delete();
	}

	virtual void STDMETHODCALLTYPE Execute(_In_ PCSTR Command)
	{
		m_capture.Execute(Command);
	}

	virtual const char* STDMETHODCALLTYPE GetTextNonNull(void)
	{
		return m_capture.GetTextNonNull();
	}

private:
	ULONG m_ref;
	ExtCaptureOutputA m_capture;
};

class CaptureOutputWImpl: public CaptureOutputW
{
public:
	CaptureOutputWImpl(): m_ref(1)
	{
	}

	virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void **ppvObject)
	{
		return E_NOTIMPL;
	}

	virtual ULONG STDMETHODCALLTYPE AddRef(void)
	{
		return InterlockedIncrement(&m_ref);
	}

	virtual ULONG STDMETHODCALLTYPE Release(void)
	{
		ULONG ret = InterlockedDecrement(&m_ref);
		if (ret == 0)
			delete this;
		return ret;
	}

	virtual void STDMETHODCALLTYPE Start(void)
	{
		capture.Start();
	}

	virtual void STDMETHODCALLTYPE Stop(void)
	{
		capture.Stop();
	}

	virtual void STDMETHODCALLTYPE Delete(void)
	{
		capture.Delete();
	}

	virtual void STDMETHODCALLTYPE Execute(_In_ PCSTR Command)
	{
		capture.Execute(Command);
	}

	virtual const WCHAR* STDMETHODCALLTYPE GetTextNonNull(void)
	{
		return capture.GetTextNonNull();
	}

private:
	ULONG m_ref;
	ExtCaptureOutputW capture;
};


//////////////////////////////////////////////////////////////////////////
// Call temporary data holder
class CallDataHolder
{
public:
	CallDataHolder():
	  m_Client6("The extension requires IDebugClient6."),
	  m_Client7("The extension requires IDebugClient7."),
	  m_Control7("The extension requires IDebugControl7."),
	  m_Symbols4("The extension requires IDebugSymbols4."),
	  m_Symbols5("The extension requires IDebugSymbols5.")
	{
	}

	~CallDataHolder()
	{
		EXT_RELEASE(m_Client6);
		EXT_RELEASE(m_Client7);
		EXT_RELEASE(m_Control7);
		EXT_RELEASE(m_Symbols4);
		EXT_RELEASE(m_Symbols5);
	}

	friend class WindbgExtImpl;
private:
#define OPT_IF(_If, _Member) \
if ((status = ext.m_Client->QueryInterface(__uuidof(_If), (PVOID*)&_Member)) != S_OK) \
{ \
	_Member.Set(NULL); \
}

	int Init(ExtExtension& ext)
	{
		HRESULT status;
		OPT_IF(IDebugClient6, m_Client6);
		OPT_IF(IDebugClient7, m_Client7);
		OPT_IF(IDebugControl7, m_Control7);
		OPT_IF(IDebugSymbols4, m_Symbols4);
		OPT_IF(IDebugSymbols5, m_Symbols5);
		return 0;
	}

	ExtCheckedPointer<IDebugClient6> m_Client6;
	ExtCheckedPointer<IDebugClient7> m_Client7;
	ExtCheckedPointer<IDebugControl7> m_Control7;
	ExtCheckedPointer<IDebugSymbols4> m_Symbols4;
	ExtCheckedPointer<IDebugSymbols5> m_Symbols5;
};


//////////////////////////////////////////////////////////////////////////
// Windbg interface
class WindbgExtImpl: public WindbgExt
{
public:
	WindbgExtImpl(ExtExtension& _ext): ext(_ext), call(NULL), dml(true)
	{
	}

	int InitCall(CallDataHolder* holder)
	{
		if (!holder)
			return 1;
		if (call)
			return 2;
		call = holder;
		return call->Init(ext);
	}

	void DeinitCall()
	{
		call = NULL;
	}

	virtual ULONG STDMETHODCALLTYPE GetVersion()
	{
		return WINDBGEXT_VERSION;
	}

	virtual const char* STDMETHODCALLTYPE GetVersionStr()
	{
		return WINDBGEXT_VERSION_STR;
	}

	// return string with comma separated list of available interfaces
	virtual const char* STDMETHODCALLTYPE GetSupportedInterfaces()
	{
		static std::string supportedIfaces;
		if (supportedIfaces.empty())
		{
			supportedIfaces = "IDebugAdvanced, IDebugClient, IDebugControl, IDebugDataSpaces, IDebugRegisters, IDebugSymbols, IDebugSystemObjects";

#define CHECK_IFACE_AVAILABLE(namePart) \
if (ext.m_ ## namePart.IsSet()) \
	supportedIfaces += ", IDebug" #namePart

#define CHECK_IFACE_AVAILABLE2(namePart, ifaceNamePart) \
if (ext.m_ ## namePart.IsSet()) \
	supportedIfaces += ", IDebug" #ifaceNamePart

#define CHECK_IFACE_AVAILABLE_FROM_HOLDER(namePart) \
if (call->m_ ## namePart.IsSet()) \
	supportedIfaces += ", IDebug" #namePart

			CHECK_IFACE_AVAILABLE(Advanced2);
			CHECK_IFACE_AVAILABLE(Advanced3);
			CHECK_IFACE_AVAILABLE(Client2);
			CHECK_IFACE_AVAILABLE(Client3);
			CHECK_IFACE_AVAILABLE(Client4);
			CHECK_IFACE_AVAILABLE(Client5);
			CHECK_IFACE_AVAILABLE_FROM_HOLDER(Client6);
			CHECK_IFACE_AVAILABLE_FROM_HOLDER(Client7);
			CHECK_IFACE_AVAILABLE(Control2);
			CHECK_IFACE_AVAILABLE(Control3);
			CHECK_IFACE_AVAILABLE(Control4);
			CHECK_IFACE_AVAILABLE(Control5);
			CHECK_IFACE_AVAILABLE(Control6);
			CHECK_IFACE_AVAILABLE_FROM_HOLDER(Control7);
			CHECK_IFACE_AVAILABLE2(Data2, DataSpaces2);
			CHECK_IFACE_AVAILABLE2(Data3, DataSpaces3);
			CHECK_IFACE_AVAILABLE2(Data4, DataSpaces4);
			CHECK_IFACE_AVAILABLE(Registers2);
			CHECK_IFACE_AVAILABLE(Symbols2);
			CHECK_IFACE_AVAILABLE(Symbols3);
			CHECK_IFACE_AVAILABLE_FROM_HOLDER(Symbols4);
			CHECK_IFACE_AVAILABLE_FROM_HOLDER(Symbols5);
			CHECK_IFACE_AVAILABLE2(System2, SystemObjects2);
			CHECK_IFACE_AVAILABLE2(System3, SystemObjects3);
			CHECK_IFACE_AVAILABLE2(System4, SystemObjects4);
		}
		return supportedIfaces.c_str();
	}

	// take list of interfaces required by script and return true if all are available and false otherwise
	virtual BOOL STDMETHODCALLTYPE CheckAvailableInterfaces(const char* ifaceList)
	{
		if (!ifaceList)
			return TRUE;

		std::string requiredHolder(ifaceList);
		char *required = const_cast<char*>(requiredHolder.c_str());
		const char *available = GetSupportedInterfaces();
		for(char *current = required, *currentEnd = NULL;
			current != NULL;
			current = currentEnd ? currentEnd + 1 : NULL)
		{
			currentEnd = strchr(current, ',');
			if (current == currentEnd)
				continue;
			if (currentEnd != NULL)
				*currentEnd = 0;

			if (strstr(available, current) == NULL)
				return FALSE;
		}

		return TRUE;
	}

	// throw an exception if these interfaces aren't available
	virtual const NoThrowException* STDMETHODCALLTYPE RequireInterfaces(const char* ifaceList)
	{
		if (!CheckAvailableInterfaces(ifaceList))
		{
			static const NoThrowException NoRequiredInterfaces = {"Some of the required interfaces aren't available"};
			return &NoRequiredInterfaces;
		}
		return NULL;
	}

#define TRY_IFACE_BEGIN()

#define TRY_IFACE(field, type) \
if (ext.m_ ## field.IsSet()) \
	return reinterpret_cast<type*>(ext.m_ ## field.Get());

#define TRY_IFACE_FROM_HOLDER(field, type) \
if (call->m_ ## field.IsSet()) \
	return reinterpret_cast<type*>(call->m_ ## field.Get());

#define TRY_IFACE_END() \
return NULL;

	virtual IDebugAdvanced3* STDMETHODCALLTYPE ExtAdvanced()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE(Advanced3, IDebugAdvanced3)
			TRY_IFACE(Advanced2, IDebugAdvanced3)
			TRY_IFACE(Advanced, IDebugAdvanced3)
		TRY_IFACE_END()
	}

	virtual IDebugClient7* STDMETHODCALLTYPE ExtClient()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE_FROM_HOLDER(Client7, IDebugClient7)
			TRY_IFACE_FROM_HOLDER(Client6, IDebugClient7)
			TRY_IFACE(Client5, IDebugClient7)
			TRY_IFACE(Client4, IDebugClient7)
			TRY_IFACE(Client3, IDebugClient7)
			TRY_IFACE(Client2, IDebugClient7)
			TRY_IFACE(Client, IDebugClient7)
		TRY_IFACE_END()
	}

	virtual IDebugControl7* STDMETHODCALLTYPE ExtControl()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE_FROM_HOLDER(Control7, IDebugControl7)
			TRY_IFACE(Control6, IDebugControl7)
			TRY_IFACE(Control5, IDebugControl7)
			TRY_IFACE(Control4, IDebugControl7)
			TRY_IFACE(Control3, IDebugControl7)
			TRY_IFACE(Control2, IDebugControl7)
			TRY_IFACE(Control, IDebugControl7)
		TRY_IFACE_END()
	}

	virtual IDebugDataSpaces4* STDMETHODCALLTYPE ExtData()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE(Data4, IDebugDataSpaces4)
			TRY_IFACE(Data3, IDebugDataSpaces4)
			TRY_IFACE(Data2, IDebugDataSpaces4)
			TRY_IFACE(Data, IDebugDataSpaces4)
		TRY_IFACE_END()
	}

	virtual IDebugRegisters2* STDMETHODCALLTYPE ExtRegisters()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE(Registers2, IDebugRegisters2)
			TRY_IFACE(Registers, IDebugRegisters2)
		TRY_IFACE_END()
	}

	virtual IDebugSymbols5* STDMETHODCALLTYPE ExtSymbols()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE_FROM_HOLDER(Symbols5, IDebugSymbols5)
			TRY_IFACE_FROM_HOLDER(Symbols4, IDebugSymbols5)
			TRY_IFACE(Symbols3, IDebugSymbols5)
			TRY_IFACE(Symbols2, IDebugSymbols5)
			TRY_IFACE(Symbols, IDebugSymbols5)
		TRY_IFACE_END()
	}

	virtual IDebugSystemObjects4* STDMETHODCALLTYPE ExtSystem()
	{
		TRY_IFACE_BEGIN()
			TRY_IFACE(System4, IDebugSystemObjects4)
			TRY_IFACE(System3, IDebugSystemObjects4)
			TRY_IFACE(System2, IDebugSystemObjects4)
			TRY_IFACE(System, IDebugSystemObjects4)
		TRY_IFACE_END()
	}

	virtual void* STDMETHODCALLTYPE ExtDummy1()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy2()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy3()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy4()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy5()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy6()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy7()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy8()
	{
		return NULL;
	}

	virtual void* STDMETHODCALLTYPE ExtDummy9()
	{
		return NULL;
	}

	virtual ULONG STDMETHODCALLTYPE Get_OutputWidth()
	{
		return ext.m_OutputWidth;
	}

	virtual ULONG STDMETHODCALLTYPE Get_ActualMachine()
	{
		return ext.m_ActualMachine;
	}

	virtual ULONG STDMETHODCALLTYPE Get_Machine()
	{
		return ext.m_Machine;
	}

	virtual ULONG STDMETHODCALLTYPE Get_PageSize()
	{
		return ext.m_PageSize;
	}

	virtual ULONG STDMETHODCALLTYPE Get_PtrSize()
	{
		return ext.m_PtrSize;
	}

	virtual ULONG STDMETHODCALLTYPE Get_NumProcessors()
	{
		return ext.m_NumProcessors;
	}

	virtual ULONG64 STDMETHODCALLTYPE Get_OffsetMask()
	{
		return ext.m_OffsetMask;
	}

	virtual ULONG STDMETHODCALLTYPE Get_DebuggeeClass()
	{
		return ext.m_DebuggeeClass;
	}

	virtual ULONG STDMETHODCALLTYPE Get_DebuggeeQual()
	{
		return ext.m_DebuggeeQual;
	}

	virtual ULONG STDMETHODCALLTYPE Get_DumpFormatFlags()
	{
		return ext.m_DumpFormatFlags;
	}

	virtual bool STDMETHODCALLTYPE Get_IsRemote()
	{
		return ext.m_IsRemote;
	}

	virtual bool STDMETHODCALLTYPE Get_OutCallbacksDmlAware()
	{
		return ext.m_OutCallbacksDmlAware;
	}

	virtual bool STDMETHODCALLTYPE IsUserMode(void)
	{
		return ext.IsUserMode();
	}

	virtual bool STDMETHODCALLTYPE IsKernelMode(void)
	{
		return ext.IsKernelMode();
	}

	virtual bool STDMETHODCALLTYPE IsLiveLocalUser(void)
	{
		return ext.IsLiveLocalUser();
	}

	virtual bool STDMETHODCALLTYPE IsMachine32(ULONG Machine)
	{
		return ext.IsMachine32(Machine);
	}

	virtual bool STDMETHODCALLTYPE IsCurMachine32(void)
	{
		return ext.IsCurMachine32();
	}

	virtual bool STDMETHODCALLTYPE IsMachine64(ULONG Machine)
	{
		return ext.IsMachine64(Machine);
	}

	virtual bool STDMETHODCALLTYPE IsCurMachine64(void)
	{
		return ext.IsCurMachine64();
	}

	virtual bool STDMETHODCALLTYPE Is32On64(void)
	{
		return ext.Is32On64();
	}

	virtual bool STDMETHODCALLTYPE CanQueryVirtual(void)
	{
		return ext.CanQueryVirtual();
	}

	virtual bool STDMETHODCALLTYPE HasFullMemBasic(void)
	{
		return ext.HasFullMemBasic();
	}

	virtual bool STDMETHODCALLTYPE IsExtensionRemote(void)
	{
		return ext.IsExtensionRemote();
	}

	virtual bool STDMETHODCALLTYPE AreOutputCallbacksDmlAware(void)
	{
		return ext.AreOutputCallbacksDmlAware();
	}

	virtual const NoThrowException* STDMETHODCALLTYPE RequireUserMode(void)
	{
		try
		{
			ext.RequireUserMode();
			return NULL;
		} catch(...)
		{
		}

		static const NoThrowException error = {"Not in user mode"};
		return &error;
	}

	virtual const NoThrowException* STDMETHODCALLTYPE RequireKernelMode(void)
	{
		try
		{
			ext.RequireKernelMode();
			return NULL;
		} catch(...)
		{
		}

		static const NoThrowException error = {"Not in kernel mode"};
		return &error;
	}

	virtual ULONG STDMETHODCALLTYPE Input(PSTR Str, ULONG StrSize)
	{
		ULONG ResSize;
		if (ext.m_Control->Input(Str, StrSize, &ResSize) == S_OK)
		{
//			strcat_s(Str, StrSize, "\n");
			return static_cast<ULONG>(strlen(Str));
		}
		return 0;
	}

	virtual ULONG STDMETHODCALLTYPE InputW(PWSTR Str, ULONG StrSize)
	{
		ULONG ResSize;
		if (ext.m_Control4->InputWide(Str, StrSize, &ResSize) == S_OK)
		{
//			wcscat_s(Str, StrSize, L"\n");
			return static_cast<ULONG>(wcslen(Str));
		}
		return 0;
	}

	virtual ULONG STDMETHODCALLTYPE Get_OutMask()
	{
		return ext.m_OutMask;
	}

	virtual void STDMETHODCALLTYPE Out(PCSTR Str)
	{
		if (dml)
			ext.Dml("%s", Str);
		else
			ext.Out("%s", Str);
	}

	virtual void STDMETHODCALLTYPE OutW(PCWSTR Str)
	{
		if (dml)
			ext.Dml(L"%s", Str);
		else
			ext.Out(L"%s", Str);
	}

	virtual void STDMETHODCALLTYPE Warn(PCSTR Str)
	{
		if (dml)
			ext.DmlWarn("%s", Str);
		else
			ext.Warn("%s", Str);
	}

	virtual void STDMETHODCALLTYPE WarnW(PCWSTR Str)
	{
		if (dml)
			ext.DmlWarn(L"%s", Str);
		else
			ext.Warn(L"%s", Str);
	}

	virtual void STDMETHODCALLTYPE Err(PCSTR Str)
	{
		if (dml)
			ext.DmlErr("%s", Str);
		else
			ext.Err("%s", Str);
	}

	virtual void STDMETHODCALLTYPE ErrW(PCWSTR Str)
	{
		if (dml)
			ext.DmlErr(L"%s", Str);
		else
			ext.Err(L"%s", Str);
	}

	virtual void STDMETHODCALLTYPE Verb(PCSTR Str)
	{
		if (dml)
			ext.DmlVerb("%s", Str);
		else
			ext.Verb("%s", Str);
	}

	virtual void STDMETHODCALLTYPE VerbW(PCWSTR Str)
	{
		if (dml)
			ext.DmlVerb(L"%s", Str);
		else
			ext.Verb(L"%s", Str);
	}

	virtual void STDMETHODCALLTYPE DmlCmdLink(PCSTR Text, PCSTR Cmd)
	{
		ext.DmlCmdLink(Text, Cmd);
	}

	virtual void STDMETHODCALLTYPE DmlCmdExec(PCSTR Text, PCSTR Cmd)
	{
		ext.DmlCmdExec(Text, Cmd);
	}

	virtual void STDMETHODCALLTYPE RefreshOutputCallbackFlags(void)
	{
		ext.RefreshOutputCallbackFlags();
	}

	virtual void STDMETHODCALLTYPE EnableDml(bool Enable)
	{
		dml = Enable;
	}

	virtual void STDMETHODCALLTYPE SetCallStatus(HRESULT Status)
	{
		ext.SetCallStatus(Status);
	}

	virtual ULONG STDMETHODCALLTYPE GetCachedSymbolTypeId(PULONG64 Cookie, PCSTR Symbol, PULONG64 ModBase)
	{
		return ext.GetCachedSymbolTypeId(Cookie, Symbol, ModBase);
	}

	virtual ULONG STDMETHODCALLTYPE GetCachedFieldOffset(__inout PULONG64 Cookie, __in PCSTR Type, __in PCSTR Field, __out_opt PULONG64 ModBase = NULL, __out_opt PULONG TypeId = NULL)
	{
		return ext.GetCachedFieldOffset(Cookie, Type, Field, ModBase, TypeId);
	}

	virtual bool STDMETHODCALLTYPE GetCachedSymbolInfo(__in ULONG64 Cookie, __out PDEBUG_CACHED_SYMBOL_INFO Info)
	{
		return ext.GetCachedSymbolInfo(Cookie, Info);
	}

	virtual bool STDMETHODCALLTYPE AddCachedSymbolInfo(__in PDEBUG_CACHED_SYMBOL_INFO Info, __out PULONG64 Cookie)
	{
		return ext.AddCachedSymbolInfo(Info, false, Cookie);
	}

	virtual void STDMETHODCALLTYPE GetModuleImagehlpInfo(__in ULONG64 ModBase, __out struct _IMAGEHLP_MODULEW64* Info)
	{
		return ext.GetModuleImagehlpInfo(ModBase, Info);
	}

	virtual bool STDMETHODCALLTYPE ModuleHasGlobalSymbols(__in ULONG64 ModBase)
	{
		return ext.ModuleHasGlobalSymbols(ModBase);
	}

	virtual bool STDMETHODCALLTYPE ModuleHasTypeInfo(__in ULONG64 ModBase)
	{
		return ext.ModuleHasTypeInfo(ModBase);
	}

	virtual PCSTR STDMETHODCALLTYPE GetExpr64(__in PCSTR Str, __in bool Signed, __in ULONG64 Limit, __out PULONG64 Val)
	{
		try
		{
			return ext.GetExpr64(Str, Signed, Limit, Val);
		} catch(...)
		{
			Err("Error: exception in GetExpr64\n");
			return NULL;
		}
	}

	virtual PCSTR STDMETHODCALLTYPE GetExprU64(__in PCSTR Str, __in ULONG64 Limit, __out PULONG64 Val)
	{
		try
		{
			return ext.GetExprU64(Str, Limit, Val);
		} catch(...)
		{
			Err("Error: exception in GetExprU64\n");
			return NULL;
		}
	}

	virtual PCSTR STDMETHODCALLTYPE GetExprS64(__in PCSTR Str, __in LONG64 Limit, __out PLONG64 Val)
	{
		try
		{
			return ext.GetExprS64(Str, Limit, Val);
		} catch(...)
		{
			Err("Error: exception in GetExprS64\n");
			return NULL;
		}
	}

	virtual CaptureOutputA* STDMETHODCALLTYPE CreateCaptureOutputA()
	{
		return new CaptureOutputAImpl();
	}

	virtual CaptureOutputW* STDMETHODCALLTYPE CreateCaptureOutputW()
	{
		return new CaptureOutputWImpl();
	}

private:
	ExtExtension& ext;
	CallDataHolder* call;
	bool dml;
};

#endif // __WINDBGEXTIMPL_H__
